import os
import os.path as osp
import argparse
import torch
import tensorflow as tf
import numpy as np
from utils import get_test_dataset


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--tf_path', type=str, default=osp.join('data', 'qat.pb'))
    parser.add_argument('--tflite_path', type=str, default=osp.join('data', 'qat.tflite'))
    parser.add_argument('--dataset', type=str, choices=['mnist', 'cifar10'], default='mnist')
    args = parser.parse_args()

    def representative_dataset_gen():
        _, test_loader = get_test_dataset(args.dataset)
        for image, _ in test_loader:
            image = np.ascontiguousarray(image.permute(0, 2, 3, 1).numpy())
            yield [image]

    converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(args.tf_path, ['input'], ['output'])
    converter.optimizations = [tf.compat.v1.lite.Optimize.DEFAULT]
    converter.inference_input_type = tf.uint8
    converter.inference_output_type = tf.uint8
    converter.quantized_input_stats = {'input': (128, 128)}
    converter.target_spec.supported_ops = [tf.compat.v1.lite.OpsSet.TFLITE_BUILTINS_INT8]
    converter.representative_dataset = representative_dataset_gen

    tflite_quant_model = converter.convert()
    with open(args.tflite_path, 'wb') as w:
        w.write(tflite_quant_model)
    print('done.')


if __name__ == '__main__':
    main()
